#include "ExecutionEngine/cuda_runtime_api_def.h"
#include "core/container.h"
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <ostream>
#include <string>
#include <iomanip>
#include <fstream>
#include <functional>

extern "C" unsigned int __sstcudaRegisterFatBinary() {
  const char *work_dir = getenv("WAFER_WORKSPACE");
  char file_name[256];
  if (work_dir != NULL) {
    snprintf(file_name, sizeof(file_name), "%s/wafer_gpu.out", work_dir);
  } else {
    snprintf(file_name, sizeof(file_name), "./wafer_gpu.out");
  }

  printf("%s", file_name);

  unsigned int fatbin_handle = __cudaRegisterFatBinary(file_name);
  return fatbin_handle;
}

extern "C"
void __sstcudaRegisterFunction(uint64_t fatCubinHandle,
                               uint64_t hostFun){
  char deviceFun[256] = "CUDA_kernel";
  __cudaRegisterFunction(fatCubinHandle, hostFun, deviceFun);
}

extern "C"
void sstcudaConfigureCall(unsigned int blockX, unsigned int blockY, unsigned int blockZ,
                          unsigned int gridX, unsigned int gridY, unsigned int gridZ,
                          uint64_t sharedMem){
  dim3 gridDim = {gridX, gridY, gridZ};
  dim3 blockDim = {blockX, blockY, blockZ};
  cudaConfigureCall(gridDim, blockDim, sharedMem);
}

extern "C"
void sstcudaLaunch(uint64_t func){
  cudaLaunch(func);
}

extern "C"
void* sstcudaMalloc(size_t size){
  void *data;
  cudaMalloc((void**)&data, size * sizeof(unsigned char));
  return (void *)data;
}

extern "C"
void sstcudaMemcpy(void *dst, void *src, size_t sizeBytes, bool kind){
  if (!kind) {
    cudaMemcpy(reinterpret_cast<uint64_t>(dst) & 0xFFFFFFFF, reinterpret_cast<uint64_t>(src),
               sizeBytes * sizeof(unsigned char), cudaMemcpyHostToDevice);
  } else {
    cudaMemcpy(reinterpret_cast<uint64_t>(dst), reinterpret_cast<uint64_t>(src) & 0xFFFFFFFF,
               sizeBytes * sizeof(unsigned char), cudaMemcpyDeviceToHost);
  }
}

extern "C"
void sstSetupIntArgument(uint64_t integer, uint64_t offset){
  cudaSetupArgument((void *) integer, 8, offset);
}

extern "C"
void sstSetupInt8Argument(uint8_t integer, uint64_t offset){
  cudaSetupArgument((void *) integer, 8, offset);
}

extern "C"
void sstSetupInt32Argument(uint32_t integer, uint64_t offset){
  cudaSetupArgument((void *) integer, 8, offset);
}

extern "C"
void sstSetupMemrefArgument(void *allocated, void *aligned, intptr_t offset, intptr_t size, intptr_t stride, uint64_t argOffset){
  cudaSetupArgument((void *) allocated, 8, argOffset);
  cudaSetupArgument((void *) aligned, 8, argOffset + 8);
  cudaSetupArgument((void *) offset, 8, argOffset + 16);
  cudaSetupArgument((void *) size, 8, argOffset + 24);
  cudaSetupArgument((void *) stride, 8, argOffset + 32);
}

extern "C"
uint64_t sstGetCpuId(){
  return (uint64_t)getCPUID();
}

extern "C"
void sstMsgSend(int src, int dst, void *buf_p, int count){
  ariel_msg_send((uint32_t)(src<<8), (uint32_t)(dst<<8), (void *)buf_p, (uint64_t)(count), 1);
}

extern "C" void sstMsgRecv(int src, int dst, void *buf_p, int count) {
  ariel_msg_recv((uint32_t)(src << 8), (uint32_t)(dst << 8), (void *)buf_p,
                 (uint64_t)(count), 1, 1);
}

extern "C" void _mlir_ciface_printScheduleInfo(unsigned int group_id,
                                               MemRef<uint8_t, 1>* data,
                                               MemRef<int32_t, 1>* devices,
                                               MemRef<float, 1>* ratios,
                                               unsigned int message_length,
                                               int operation_type) {
  // Generate filename like "group_0.json"
  std::string filename = "group_" + std::to_string(group_id) + ".json";
  std::ofstream File(filename, std::ios::app);
  File << std::unitbuf; // Refresh the buffer before writed into file.
  if (!File.is_open()) {
    perror("Failed to open log file.");
    return;
  }

  // Prepare all values
  auto data_size = data->getSizes()[0];
  auto device_count = devices->getSizes()[0];
  auto ratios_count = ratios->getSizes()[0];
  assert(ratios_count == device_count);

  // Functions for data formatting
  std::function<void(uint8_t)> add_to_buffer;
  HexFormatter add_hex_to_buffer(File);
  DecimalFormatter add_decimal_to_buffer(File);

  // Determine operation type specifics
  bool is_input = false;
  bool is_multi = false;
  bool is_extra = false;
  std::string crypto_name = "";
  switch (operation_type) {
      // is_input ?
      case OP_TYPE::INPUT_MULTI_SHA1:
      case OP_TYPE::INPUT_MULTI_MD5:
      case OP_TYPE::INPUT_MULTI_SM3:
      case OP_TYPE::INPUT_MULTI_AES:
      case OP_TYPE::INPUT_CRACK_SHA1:
      case OP_TYPE::INPUT_CRACK_MD5:
      case OP_TYPE::INPUT_CRACK_SM3:
          is_input = true;
          add_to_buffer = std::function<void(uint8_t)>(add_decimal_to_buffer);
          break;
      case OP_TYPE::INPUT_CRACK_AES:
          is_input = true;
          add_to_buffer = std::function<void(uint8_t)>(add_hex_to_buffer);
          break;
      // is_multi ?
      case OP_TYPE::OUTPUT_MULTI_SHA1:
      case OP_TYPE::OUTPUT_MULTI_MD5:
      case OP_TYPE::OUTPUT_MULTI_SM3:
      case OP_TYPE::OUTPUT_MULTI_AES:
          is_input = false;
          is_multi = true;
          add_to_buffer = std::function<void(uint8_t)>(add_hex_to_buffer);
          break;
      case OP_TYPE::OUTPUT_CRACK_AES:
          is_input = false;
          is_multi = true;
          add_to_buffer = std::function<void(uint8_t)>(add_decimal_to_buffer);
          break;
      // is_extra ?
      case OP_TYPE::EXTRA_CRACK_SHA1:
      case OP_TYPE::EXTRA_CRACK_MD5:
      case OP_TYPE::EXTRA_CRACK_SM3:
          is_input = false;
          is_multi = false;
          is_extra = true;
          add_to_buffer = std::function<void(uint8_t)>(add_hex_to_buffer);
          break;
      case OP_TYPE::EXTRA_MULTI_CRACK_AES:
          is_input = false;
          is_multi = false;
          is_extra = true;
          add_to_buffer = std::function<void(uint8_t)>(add_decimal_to_buffer);
          break;
      // else
      case OP_TYPE::OUTPUT_CRACK_SHA1:
      case OP_TYPE::OUTPUT_CRACK_MD5:
      case OP_TYPE::OUTPUT_CRACK_SM3:
          is_input = false; 
          is_multi = false;
          is_extra = false;
          add_to_buffer = std::function<void(uint8_t)>(add_decimal_to_buffer);
          break;
      default: 
          std::cerr << "Unknown operation type" << std::endl;
          break;
  }
  switch (operation_type % CryptoTypesNum) {
      case 0: crypto_name = "SHA1"; break;
      case 1: crypto_name = "MD5";  break;
      case 2: crypto_name = "SM3";  break;
      case 3: crypto_name = "AES";  break;
      default: 
          std::cerr << "Unknown crytpo operation" << std::endl;
          break;
  }

  // Indentation for JSON formatting
  std::string indentation = "";

  // Write values into the File
  // For input operations
  if (is_input) {
    File << "{\n";
    indentation.push_back('\t');
    File << indentation << ADD_QUOTES("group") << ": " << group_id << ",\n";
    File << indentation << ADD_QUOTES("compileType") << ": " << ADD_QUOTES("crypto") << ",\n";
    File << indentation << ADD_QUOTES("opName") << ": " << ADD_QUOTES(crypto_name) << ",\n";
    File << indentation << ADD_QUOTES("scheduleInfo") << ": [\n";
    indentation.push_back('\t');

    int start = 0, end = 0;
    for (int i = 0; i < device_count; i++) {
      File.flush();/*If File.flush is not executed here, it will lead to a memory leak*/
      auto device = (*devices)[i];
      std::string device_str;
      switch (device) {
        case 0: device_str = "CPU"; break;
        case 1: device_str = "GPU"; break;
        case 2: device_str = "CGRA"; break;
        case 3: device_str = "FPGA"; break;
        default: 
        std::cerr << "Unknown device when printScheduleInfo: " << device_str << std::endl;
        assert(false);
      }
      device_str += std::to_string(group_id);
      auto ratio = (*ratios)[i];
      end = (i == device_count - 1) ? data_size : start + ratio * data_size;

      File << indentation << "{\n";
      indentation.push_back('\t');
      File << indentation << ADD_QUOTES("device") << ": " << ADD_QUOTES(device_str) << ",\n";
      File << indentation << ADD_QUOTES("ratio") << ": " << ratio * 100 << ",\n";
      File << indentation << ADD_QUOTES("inputNum") << ": " << std::to_string((end - start) / message_length) << ",\n";
      File << indentation << ADD_QUOTES("inputData") << ": [\n";
      
      indentation.push_back('\t');
      for (int j = start; j < end; j += message_length) {
        File << indentation << '\"';
        File.flush();/*If File.flush is not executed here, it will lead to a memory leak*/
        for (size_t k = 0; k < message_length; ++k) {
          add_to_buffer((*data)[j + k]);
        }
        File << '\"';
        if ((size_t)j < end - message_length) 
          File << ',';
        File << '\n';
      }
      indentation.pop_back();
      File << indentation << "]\n";

      indentation.pop_back();
      File << indentation << '}';
      if (i < device_count - 1)
        File << ',';
      File << '\n';

      start = end;
    }

    indentation.pop_back();
    File << indentation << "],\n";
  }
  // For multi-output operations (MD5/SM3/SHA1/AES multi or AES crack output)
  else if (is_multi) {
    indentation.push_back('\t');
    File << indentation << ADD_QUOTES("result") << ": [\n";
    indentation.push_back('\t');
    for (int j = 0; j < data_size; j += message_length) {
      File << indentation << '\"';
      File.flush();/*If File.flush is not executed here, it will lead to a memory leak*/
      for (size_t k = 0; k < message_length; ++k) {
        add_to_buffer((*data)[j + k]);
      }
      File << '\"';
      if (j < data_size - message_length) 
        File << ',';
      File << '\n';
    }
    indentation.pop_back();
    File << indentation << "]\n}\n";
  } 
  // For target (MD5/SM3/SHA1 crack) or key (AES multi/crack)
  else if (is_extra) {
    indentation.push_back('\t');
    if (crypto_name == "AES")
      File << indentation << ADD_QUOTES("key") << ": [\n";
    else 
      File << indentation << ADD_QUOTES("target") << ": [\n";
    indentation.push_back('\t');
    File << indentation << '\"';
    for (int j = 0; j < data_size; ++j) {
      add_to_buffer((*data)[j]);
    }
    indentation.pop_back();
    File << "\"\n" << indentation << "],\n";
  }
  // For single-output operations (MD5/SM3/SHA1 crack output)
  else {
    indentation.push_back('\t');
    File << indentation << ADD_QUOTES("result") << ": [\n";
    indentation.push_back('\t');
    File << indentation << '\"';
    for (int j = 0; j < data_size; ++j) {
      add_to_buffer((*data)[j]);
    }
    indentation.pop_back();
    File << "\"\n" << indentation << "]\n}\n";
  }
}

extern "C" {
void cgra_run_sha1(MemRef<unsigned char, 1> *input,
                   MemRef<unsigned char, 1> *output) {
  std::cout << "calling cgra run sha1." << std::endl;
}

void cgra_run_md5(MemRef<unsigned char, 1> *input,
                  MemRef<unsigned char, 1> *output) {
  std::cout << "calling cgra run md5." << std::endl;
}

void cgra_run_sm3(MemRef<unsigned char, 1> *input,
                  MemRef<unsigned char, 1> *output) {
  std::cout << "calling cgra run sm3." << std::endl;
}

void cgra_run_aes(MemRef<unsigned char, 1> *input,
                  MemRef<unsigned char, 1> *output) {
  std::cout << "calling cgra run aes." << std::endl;
}
}